import torch
import os
import numpy as np
import cv2
import matplotlib.pyplot as plt
from torchvision import datasets, transforms
import torchvision
from PIL import Image
from xml.dom.minidom import parse
import utils
import transforms as T
from engine import train_one_epoch, evaluate
import xml.etree.cElementTree as ET
import collections
import pandas as pd
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
import sys

# 把要显示的数据集放在这个文件夹里
root = r'dataset/train'
# root = r'dataset\train'
img_path = os.path.join(root, r'JPEGImages')
xml_path = os.path.join(root, r'Annotations')
img_list = os.listdir(img_path)
xml_list = os.listdir(xml_path)

print(img_list)
print(xml_list)

for i in range(len(img_list)):
    print(os.path.join(root, img_list[i]))
    print(os.path.join(root, xml_list[i]))
    img = cv2.imread(os.path.join(root, 'JPEGImages', img_list[i]))

    bbox_xml_path = os.path.join(root, 'Annotations', xml_list[i])
    dom = parse(bbox_xml_path)
    data = dom.documentElement
    objects = data.getElementsByTagName('object')
    labels = []
    for object_ in objects:
        # name就是label字符串
        name = object_.getElementsByTagName('name')[0].childNodes[0].nodeValue  # 就是label

        print(name)

        # 返回的应该是一个列表，但是这里只有一个bndbox，但是仍然要用下标0来获得第一个的对象
        bndbox = object_.getElementsByTagName('bndbox')[0]
        xmin = np.float(bndbox.getElementsByTagName('xmin')[0].childNodes[0].nodeValue)
        ymin = np.float(bndbox.getElementsByTagName('ymin')[0].childNodes[0].nodeValue)
        xmax = np.float(bndbox.getElementsByTagName('xmax')[0].childNodes[0].nodeValue)
        ymax = np.float(bndbox.getElementsByTagName('ymax')[0].childNodes[0].nodeValue)
        # 列表汇总，附加一个四个数的列表
        # boxes.append([xmin, ymin, xmax, ymax])
        cv2.rectangle(img, (int(xmin), int(ymin)), (int(xmax), int(ymax)), (255, 0, 0), thickness=2)
        cv2.putText(img, name, (int(xmin), int(ymin)), fontFace=cv2.FONT_HERSHEY_SIMPLEX, fontScale=1,
                    color=(0, 255, 0))

    cv2.imshow("img", img)
    key = cv2.waitKey()
    # 按q退出
    if key == ord('q'):
        break

cv2.destroyAllWindows()
